import torch
import torch.fx
from torch.fx import Interpreter

# 定义一个简单的模型
class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 2)
        self.relu = torch.nn.ReLU()
    
    def forward(self, x):
        x = self.linear(x)
        x = self.relu(x)
        return x

# 创建模型实例
model = MyModel()

# 使用FX进行符号追踪
traced = torch.fx.symbolic_trace(model)
print("原始图:")
print(traced.graph)

# 自定义Interpreter来记录操作
class MyInterpreter(Interpreter):
    def __init__(self, module, garbage_collect_values=True):
        super().__init__(module, garbage_collect_values)
        self.node_operations = []
    
    def run(self, *args, **kwargs):
        # 在运行前清空记录
        self.node_operations = []
        return super().run(*args, **kwargs)
    
    def run_node(self, n):
        # 记录节点信息
        self.node_operations.append(f"执行节点: {n.op} {n.target}")
        
        # 调用父类方法实际执行节点
        result = super().run_node(n)
        
        # 可以在这里添加自定义处理
        if n.op == 'call_function' and n.target == torch.relu:
            self.node_operations.append("检测到ReLU操作")
        
        return result
    
    def call_function(self, target, args, kwargs):
        # 可以拦截特定函数调用
        if target == torch.relu:
            print("拦截到ReLU函数调用")
        return super().call_function(target, args, kwargs)

# 使用自定义Interpreter
interpreter = MyInterpreter(traced)

# 创建输入数据
input_data = torch.randn(1, 3)

# 运行Interpreter
output = interpreter.run(input_data)
print("\nInterpreter输出:", output)

# 打印记录的操作
print("\n执行的操作记录:")
for op in interpreter.node_operations:
    print(op)

# 另一个例子：修改图的行为
class ReplacingInterpreter(Interpreter):
    def call_function(self, target, args, kwargs):
        # 将所有ReLU替换为Sigmoid
        if target == torch.relu:
            return torch.sigmoid(args[0])
        return super().call_function(target, args, kwargs)

# 使用替换Interpreter
repl_interpreter = ReplacingInterpreter(traced)
repl_output = repl_interpreter.run(input_data)
print("\n替换ReLU后的输出:", repl_output)